import os
import argparse
import json
from openai import OpenAI

def optimize_prompt_with_deepseek(instruction, context=None):
    """
    Uses the DeepSeek API to optimize a given prompt.

    Args:
        instruction (str): The main instruction to be optimized.
        context (str, optional): Additional context for the instruction. Defaults to None.

    Returns:
        dict or str: The optimized prompt, hopefully in a JSON object, or an error/raw response.
    """
    api_key = os.getenv("DEEPSEEK_API_KEY")
    if not api_key:
        return {"error": "DEEPSEEK_API_KEY environment variable not set."}

    # Point the client to the DeepSeek API endpoint
    client = OpenAI(
        api_key=api_key,
        base_url="https://api.deepseek.com/v1"
    )

    # We strongly instruct the model to return JSON, even if it's not enforced by the API.
    system_prompt = "你是一个提示词优化专家，旨在将用户输入优化得更清晰、更有效、更具结构化。请以JSON格式返回优化结果。你的回答应该是一个完整的、可被直接解析的JSON对象，不要包含任何JSON格式之外的解释性文字或代码标记。返回格式必须是：{\"optimized_prompt\": \"优化后的提示词内容\"}"

    if context:
        user_prompt = f"请优化以下指令：\n\n### 指令:\n{instruction}\n\n### 上下文:\n{context}"
    else:
        user_prompt = f"请优化以下指令：\n\n### 指令:\n{instruction}"

    try:
        print("向 DeepSeek API 发送请求...")
        response = client.chat.completions.create(
            model="deepseek-chat",  # Use the recommended chat model
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.7,
        )
        
        print("成功接收到 API 响应。")
        response_content = response.choices[0].message.content
        
        # Try to parse the response as JSON. If it fails, return the raw text.
        try:
            json_response = json.loads(response_content)
            return json_response
        except json.JSONDecodeError:
            print("\n[警告] API未返回有效的JSON格式，将以纯文本形式展示结果。")
            return {"optimized_prompt_raw_text": response_content}


    except Exception as e:
        return {"error": f"An API error occurred: {e}"}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Optimize a prompt using the DeepSeek API.")
    parser.add_argument("--instruction", type=str, required=True, help="The instruction to be optimized.")
    parser.add_argument("--context", type=str, default=None, help="Optional context for the instruction.")
    
    args = parser.parse_args()
    
    optimized_result = optimize_prompt_with_deepseek(args.instruction, args.context)
    
    # Use json.dumps to format the output, ensuring Chinese characters are displayed correctly.
    print("\n--- 优化结果 ---\n")
    print(json.dumps(optimized_result, indent=2, ensure_ascii=False)) 